import pyaudio
import wave
import numpy as np
import rospy
from audio_common_msgs.msg import AudioData
from std_msgs.msg import String
from struct import pack
import os
import sys
import itertools
import glob
import argparse
import numpy as np
from gmm_code.features import get_feature
from array import array

import keras
import matplotlib.pyplot as plt
from keras.callbacks import LambdaCallback
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import LSTM, Flatten
from keras.optimizers import RMSprop, Adam, Nadam,Adamax
from keras.utils.data_utils import get_file
import numpy as np
import random
import sys
import io
from keras.layers import Input, Dense, LSTM, RepeatVector, Reshape, Permute,Bidirectional
from keras.models import Model
from keras.layers import Bidirectional, concatenate, Conv1D, MaxPooling1D, GlobalMaxPooling1D, Dropout, BatchNormalization
import gc
from keras.utils import plot_model, np_utils
import time
import datetime
from keras.utils import to_categorical
from keras.callbacks import EarlyStopping
import warnings
from sklearn.model_selection import train_test_split
from numpy import random, mat
import multiprocessing
import os
import time
import math
from keras.utils import plot_model
warnings.filterwarnings('ignore')
from matplotlib import pyplot as plt
import itertools
from sklearn.metrics import confusion_matrix
from keras.models import load_model
from python_speech_features import mfcc
import pandas as pd
import numpy as np
import gmm_code.utils import read_wav

class MyAudio(object):
    audio_data = []

    def __init__(self):
        self.sample_rate = rospy.get_param('audio_capture/sample_rate')
        self.depth = 1  # rospy.get_param('audio_capture/depth')
        self.channels = rospy.get_param('audio_capture/channels')
        self.pub = rospy.Publisher('recognition', String, queue_size=1, tcp_nodelay=True)
        self.audio_sub = rospy.Subscriber('/audio', AudioData, self.recode_data)
        self.audio_data = []
        print('Please begin to speak ,for 5 s')
        self.save_data()

    def record_to_file(self,data,wav_name):
        "Records from the microphone and outputs the resulting data to 'path'"
        #data = pack('<' + ('h' * len(data)), *data)
        wf = wave.open(wav_name, 'wb')
        wf.setnchannels(self.channels)
        wf.setsampwidth(2)
        wf.setframerate(self.sample_rate)
        #wf.writeframes(data)
        wf.writeframes(b''.join(data))
        #print(len(data))
        wf.close()

    def recode_data(self, msg):
        j = 0
        for i in range(160):
            aa = (msg.data[j]&0xff) | ((msg.data[j+1]&0xff)<<8)
            j += 2
            self.audio_data.append(aa)

    def save_data(self):
        rate = rospy.Rate(100)
        count = 0
        while not rospy.is_shutdown():
            count += 1
            if count > 500:
                self.aa = np.array(self.audio_data, dtype=np.uint16)
                #print(self.aa)
                #wav_name = 'out_audio.wav'
                #self.record_to_file(self.aa,wav_name )
                #fs, signal = read_wav(wav_name)
                self.predict(self.aa)
                break
            rate.sleep()

    def deal_data(self,train_date_c1s):
        train_date_c1s = train_date_c1s.astype(float)
        vgac1 = np.mean(train_date_c1s)
        varc1 = math.sqrt(np.var(train_date_c1s))
        train_date_c1s = (train_date_c1s - vgac1) / varc1
        return train_date_c1s

    def predict(self,data):
        self.model = load_model('/home/banana/bear_voice/src/mfcc_tcnn_speaker/script/model/model_0498_acc=1.0.h5')
        df_feature = get_feature(self.sample_rate, data)
        df_feature = np.nan_to_num(df_feature)
        mfcc_max = 495
        if df_feature.shape[0] >  mfcc_max :
            sert = int( df_feature.shape[0] - mfcc_max )
            for j in range(sert):
                df_feature = np.delete(df_feature, -1,axis=0)
            #zero_insert = np.array(np.zeros((sert, 13))).astype(float)
            #df_feature = np.append(df_feature, zero_insert,axis=0)
        df_feature_scale = self.deal_data(df_feature)
        df_feature_scale = np.reshape(df_feature_scale, (1, mfcc_max,13))
        print(df_feature_scale,df_feature_scale.shape)
        pred_y = self.model.predict(df_feature_scale)
        pred_label = np.argmax(pred_y, axis=1)
        if pred_label == 0:
                speaker_label = 'unknown'
        else:
                speaker_label = 'linkunling'
        output_str = "speaker:  %s" %(speaker_label)
        print(output_str)
        self.pub.publish(output_str)



def main():
    '''main'''
    rospy.init_node("node_name", anonymous=True)
    MyAudio()
    print('End')


if __name__ == '__main__':
    while True:
          main()

